import os
from typing import Optional, List

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from pydantic.v1.fields import Field
from pydantic.v1.main import BaseModel

load_dotenv()


# 1.创建模型
model = ChatOpenAI(
    model='qwen-plus',
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)


# pydantic 处理数据、验证数据、定义数据格式、序列化和反序列化、类型转换等等

class Person(BaseModel):
    """
    关于一个人的数据模型
    """
    name: Optional[str] = Field(
        default=None, description="人的名字"
    )
    hair_color: Optional[str] = Field(
        default=None, description="如果知道的话，这个人的头发颜色"
    )
    height_in_meters: Optional[str] = Field(
        default=None, description="以米为单位测量的身高高度"
    )

class ManyPerson(BaseModel):
    people: Optional[List[Person]] = Field(default=None, description="多个人")


system = """你是一个专业的提取算法。只从未结构化文本中提取相关信息。
如果你不知道要提取的属性的值，返回该属性的值为null。
"""
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        # MessagesPlaceholder("examples"),
        ("human", "{text}"),
    ]
)

text = '1米9身高的李明从对面马路走过来了，染了个新的棕色头发，好看极了！张三身高比李明矮10厘米。'

# chain = {'text': RunnablePassthrough()} | prompt | model.with_structured_output(schema=Person)
chain = {'text': RunnablePassthrough()} | prompt | model.with_structured_output(schema=ManyPerson)

resp = chain.invoke(text)

print(resp)
